load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("//bazel:tachyon.bzl", "if_gpu_is_configured")
load(
    "//bazel:tachyon_cc.bzl",
    "tachyon_cc_library",
    "tachyon_cuda_defines",
    "tachyon_cuda_unittest",
)

package(default_visibility = ["//visibility:public"])

tachyon_cc_library(
    name = "device_perf_info",
    srcs = ["device_perf_info.cc"],
    hdrs = ["device_perf_info.h"],
    deps = [
        "//tachyon:export",
        "//tachyon/base:no_destructor",
        "//tachyon/build:build_config",
        "@com_google_absl//absl/synchronization",
    ],
)

tachyon_cc_library(
    name = "gpu_cudamallocasync_allocator",
    srcs = if_cuda(["gpu_cudamallocasync_allocator.cc"]),
    hdrs = ["gpu_cudamallocasync_allocator.h"],
    defines = tachyon_cuda_defines(),
    deps = [
        ":gpu_logging",
        ":platform_device_id",
        "//tachyon/device:allocator",
    ],
)

tachyon_cc_library(
    name = "gpu_device_functions",
    hdrs = ["gpu_device_functions.h"],
    defines = tachyon_cuda_defines(),
    deps = if_cuda([
        "@local_config_cuda//cuda:cuda_headers",
    ]),
)

tachyon_cc_library(
    name = "gpu_driver",
    srcs = if_gpu_is_configured(["gpu_driver.cc"]),
    hdrs = ["gpu_driver.h"],
    deps = [
        ":gpu_types",
        "//tachyon/base:logging",
    ],
)

tachyon_cc_library(
    name = "gpu_enums",
    hdrs = ["gpu_enums.h"],
    defines = tachyon_cuda_defines(),
    deps = if_cuda([
        "@local_config_cuda//cuda:cuda_headers",
    ]),
)

tachyon_cc_library(
    name = "gpu_info",
    srcs = [
        "gpu_info.cc",
        "gpu_util.cc",
    ],
    hdrs = [
        "gpu_info.h",
        "gpu_util.h",
    ],
    deps = [
        ":device_perf_info",
        "//tachyon/base/strings:string_number_conversions",
        "//tachyon/base/time",
    ],
)

tachyon_cc_library(
    name = "gpu_init",
    srcs = ["gpu_init.cc"],
    hdrs = ["gpu_init.h"],
    defines = tachyon_cuda_defines(),
    deps = ["//tachyon/base:logging"],
)

tachyon_cc_library(
    name = "gpu_logging",
    srcs = if_gpu_is_configured(["gpu_logging.cc"]),
    hdrs = ["gpu_logging.h"],
    deps = [
        ":gpu_device_functions",
        ":gpu_enums",
        "//tachyon/base:compiler_specific",
        "//tachyon/base:logging",
    ],
)

tachyon_cc_library(
    name = "gpu_memory",
    srcs = if_gpu_is_configured(["gpu_memory.cc"]),
    hdrs = ["gpu_memory.h"],
    deps = [
        ":gpu_logging",
        "//tachyon:export",
        "//tachyon/base/numerics:checked_math",
        "@com_google_absl//absl/types:span",
    ],
)

tachyon_cc_library(
    name = "gpu_types",
    hdrs = ["gpu_types.h"],
    defines = tachyon_cuda_defines(),
    deps = if_cuda([
        "@local_config_cuda//cuda:cuda_headers",
    ]),
)

tachyon_cc_library(
    name = "platform_device_id",
    hdrs = ["platform_device_id.h"],
)

tachyon_cc_library(
    name = "scoped_event",
    srcs = if_gpu_is_configured(["scoped_event.cc"]),
    hdrs = ["scoped_event.h"],
    deps = [
        ":gpu_logging",
        "//tachyon:export",
    ],
)

tachyon_cc_library(
    name = "scoped_mem_pool",
    srcs = if_gpu_is_configured(["scoped_mem_pool.cc"]),
    hdrs = ["scoped_mem_pool.h"],
    deps = [
        ":gpu_logging",
        "//tachyon:export",
    ],
)

tachyon_cc_library(
    name = "scoped_stream",
    srcs = if_gpu_is_configured(["scoped_stream.cc"]),
    hdrs = ["scoped_stream.h"],
    deps = [
        ":gpu_logging",
        "//tachyon:export",
    ],
)

tachyon_cuda_unittest(
    name = "gpu_unittests",
    srcs = if_gpu_is_configured(["gpu_memory_unittest.cc"]),
    deps = [
        ":gpu_memory",
        ":scoped_mem_pool",
        "//tachyon/base:random",
        "//tachyon/base/containers:container_util",
    ],
)
